from langchain_core.tools import tool
from pydantic import BaseModel, Field
from typing import Dict

# 定义输入模型
class ContributionCalculationInput(BaseModel):
    individual_results: Dict[str, float] = Field(description="每个维度的计算结果")
    overall_result: float = Field(description="整体指标结果")

# 贡献度计算工具
@tool("contribution_calculation", args_schema=ContributionCalculationInput, return_direct=True)
async def contribution_calculation(individual_results: Dict[str, float], overall_result: float) -> Dict[str, float]:
    """计算每个维度对整体指标的贡献度"""
    try:
        contribution = {}
        for dimension, value in individual_results.items():
            contribution[dimension] = value / overall_result
        return contribution
    except Exception as e:
        print(f"Error calculating contribution: {e}")
        return {}
